import winreg
import re
import datetime
import os
import csv
import sys # 用于在Windows上可靠地检查管理员权限

def filetime_to_datetime(ft):
    """
    将Windows FILETIME（自1601-01-01 UTC以来100纳秒间隔的数量）
    转换为本地时间的Python datetime对象。
    """
    # FILETIME是一个64位值，表示自1601年1月1日UTC以来100纳秒间隔的数量。
    # 要转换为Unix纪元（1970年1月1日UTC），需要减去1601-01-01和1970-01-01之间100纳秒间隔的数量。
    # 这个常数是116444736000000000（秒 * 10,000,000 以获得100纳秒间隔）
    unix_epoch_diff_100ns = 116444736000000000
    
    # 计算自Unix纪元以来的微秒数
    microseconds_since_unix_epoch = (ft - unix_epoch_diff_100ns) // 10
    
    # 从Unix纪元在UTC中创建一个朴素的datetime对象
    dt_utc = datetime.datetime(1970, 1, 1, tzinfo=datetime.timezone.utc) + datetime.timedelta(microseconds=microseconds_since_unix_epoch)
    
    # 转换为本地时区
    return dt_utc.astimezone(datetime.datetime.now().astimezone().tzinfo)

def get_usb_info_from_registry():
    """
    从Windows注册表提取USB设备信息。
    主要关注USBSTOR（大容量存储）和通用USB设备。
    """
    usb_devices = {} # 使用字典存储唯一设备，键：(供应商ID, 产品ID, 序列号)

    # --- 1. USB大容量存储设备 (USBSTOR) ---
    try:
        key_path_usbstor = r"SYSTEM\CurrentControlSet\Enum\USBSTOR"
        with winreg.OpenKey(winreg.HKEY_LOCAL_MACHINE, key_path_usbstor) as hkey_usbstor:
            for i in range(winreg.QueryInfoKey(hkey_usbstor)[0]):
                try:
                    # 例如：Disk&Ven_SanDisk&Prod_Cruzer_Blade&Rev_1.00
                    subkey_name_disk = winreg.EnumKey(hkey_usbstor, i) 
                    
                    # 从磁盘子键名称中提取供应商ID和产品ID
                    vid_match = re.search(r"Ven_([0-9A-F]{4})", subkey_name_disk, re.IGNORECASE)
                    pid_match = re.search(r"Prod_([0-9A-F]{4})", subkey_name_disk, re.IGNORECASE)
                    vid = vid_match.group(1).upper() if vid_match else "Unknown"
                    pid = pid_match.group(1).upper() if pid_match else "Unknown"

                    with winreg.OpenKey(hkey_usbstor, subkey_name_disk) as hkey_disk:
                        for j in range(winreg.QueryInfoKey(hkey_disk)[0]):
                            try:
                                # 这是实例ID，通常包含序列号
                                # 例如：20044352710192F06053&0
                                subkey_name_instance = winreg.EnumKey(hkey_disk, j) 
                                
                                # 提取序列号：通常是实例ID中第一个'&'之前的部分
                                serial = subkey_name_instance.split('&')[0].upper()
                                # 如果分割结果为空字符串（例如："&&0"），则回退
                                if not serial: 
                                    serial = subkey_name_instance.upper()
                                    
                                device_key = (vid, pid, serial)
                                
                                # 如果设备信息尚未存在，则初始化
                                if device_key not in usb_devices:
                                    usb_devices[device_key] = {
                                        "vid": vid,
                                        "pid": pid,
                                        "serial_number": serial,
                                        "device_type": "USB Mass Storage",
                                        "description": "",
                                        "friendly_name": "",
                                        "manufacturer": "",
                                        "first_seen_setupapi": None,
                                        "last_seen_registry": None,
                                        "registry_path": f"{key_path_usbstor}\\{subkey_name_disk}\\{subkey_name_instance}"
                                    }
                                
                                with winreg.OpenKey(hkey_disk, subkey_name_instance) as hkey_instance:
                                    try:
                                        usb_devices[device_key]["description"] = winreg.QueryValueEx(hkey_instance, "DeviceDesc")[0]
                                    except FileNotFoundError:
                                        pass
                                    try:
                                        usb_devices[device_key]["friendly_name"] = winreg.QueryValueEx(hkey_instance, "FriendlyName")[0]
                                    except FileNotFoundError:
                                        pass
                                    try:
                                        usb_devices[device_key]["manufacturer"] = winreg.QueryValueEx(hkey_instance, "Mfg")[0]
                                    except FileNotFoundError:
                                        pass
                                    
                                    # 获取实例键的LastWriteTime作为最后一次看到的代理时间
                                    _, _, last_write_ft = winreg.QueryInfoKey(hkey_instance)
                                    last_seen_dt = filetime_to_datetime(last_write_ft)
                                    
                                    # 如果这个时间更晚或者这是新条目，则更新last_seen_registry
                                    if usb_devices[device_key]["last_seen_registry"] is None or \
                                       last_seen_dt > usb_devices[device_key]["last_seen_registry"]:
                                        usb_devices[device_key]["last_seen_registry"] = last_seen_dt

                            except WindowsError as e:
                                # print(f"读取USBSTOR实例子键 '{subkey_name_instance}' 时出错: {e}")
                                continue
                except WindowsError as e:
                    # print(f"读取USBSTOR磁盘子键 '{subkey_name_disk}' 时出错: {e}")
                    continue
    except FileNotFoundError:
        print(f"注册表键未找到: {key_path_usbstor} (如果没有连接USB存储设备，这是正常的)")
    except WindowsError as e:
        print(f"访问注册表键 {key_path_usbstor} 时出错: {e}")

    # --- 2. 通用USB设备 (USB) ---
    try:
        key_path_usb = r"SYSTEM\CurrentControlSet\Enum\USB"
        with winreg.OpenKey(winreg.HKEY_LOCAL_MACHINE, key_path_usb) as hkey_usb_root:
            for i in range(winreg.QueryInfoKey(hkey_usb_root)[0]):
                try:
                    # 例如：VID_04F2&PID_B107
                    subkey_name_vidpid = winreg.EnumKey(hkey_usb_root, i) 
                    vid_match = re.search(r"VID_([0-9A-F]{4})", subkey_name_vidpid, re.IGNORECASE)
                    pid_match = re.search(r"PID_([0-9A-F]{4})", subkey_name_vidpid, re.IGNORECASE)
                    vid = vid_match.group(1).upper() if vid_match else "Unknown"
                    pid = pid_match.group(1).upper() if pid_match else "Unknown"

                    with winreg.OpenKey(hkey_usb_root, subkey_name_vidpid) as hkey_vidpid:
                        for j in range(winreg.QueryInfoKey(hkey_vidpid)[0]):
                            try:
                                # 这通常是序列号
                                subkey_name_instance = winreg.EnumKey(hkey_vidpid, j) 
                                # 对于通用USB，这通常就是序列号本身
                                serial = subkey_name_instance.upper() 

                                device_key = (vid, pid, serial)
                                
                                if device_key not in usb_devices:
                                    usb_devices[device_key] = {
                                        "vid": vid,
                                        "pid": pid,
                                        "serial_number": serial,
                                        "device_type": "General USB Device",
                                        "description": "",
                                        "friendly_name": "",
                                        "manufacturer": "",
                                        "first_seen_setupapi": None,
                                        "last_seen_registry": None,
                                        "registry_path": f"{key_path_usb}\\{subkey_name_vidpid}\\{subkey_name_instance}"
                                    }
                                
                                with winreg.OpenKey(hkey_vidpid, subkey_name_instance) as hkey_instance:
                                    try:
                                        usb_devices[device_key]["description"] = winreg.QueryValueEx(hkey_instance, "DeviceDesc")[0]
                                    except FileNotFoundError:
                                        pass
                                    try:
                                        usb_devices[device_key]["friendly_name"] = winreg.QueryValueEx(hkey_instance, "FriendlyName")[0]
                                    except FileNotFoundError:
                                        pass
                                    try:
                                        usb_devices[device_key]["manufacturer"] = winreg.QueryValueEx(hkey_instance, "Mfg")[0]
                                    except FileNotFoundError:
                                        pass
                                    
                                    _, _, last_write_ft = winreg.QueryInfoKey(hkey_instance)
                                    last_seen_dt = filetime_to_datetime(last_write_ft)
                                    
                                    # 如果这个时间更晚或者这是新条目，则更新last_seen_registry
                                    if usb_devices[device_key]["last_seen_registry"] is None or \
                                       last_seen_dt > usb_devices[device_key]["last_seen_registry"]:
                                        usb_devices[device_key]["last_seen_registry"] = last_seen_dt
                                    
                                    # 如果之前没有设备类型，或者类型是“从SetupAPI日志”，则更新为“通用USB设备”
                                    # （此逻辑确保如果USBSTOR找到该设备，则USBSTOR类型优先）
                                    if not usb_devices[device_key].get("device_type") or usb_devices[device_key]["device_type"] == "从SetupAPI日志":
                                        usb_devices[device_key]["device_type"] = "General USB Device"

                            except WindowsError as e:
                                # print(f"读取USB实例子键 '{subkey_name_instance}' 时出错: {e}")
                                continue
                except WindowsError as e:
                    # print(f"读取USB VID_PID子键 '{subkey_name_vidpid}' 时出错: {e}")
                    continue
    except FileNotFoundError:
        print(f"注册表键未找到: {key_path_usb} (如果没有连接通用USB设备，这是正常的)")
    except WindowsError as e:
        print(f"访问注册表键 {key_path_usb} 时出错: {e}")
            
    return usb_devices


def get_first_install_times_from_setupapi(devices_info):
    """
    解析setupapi.dev.log以获取USB设备安装时间戳。
    更新devices_info字典中的'first_seen_setupapi'字段。
    """
    log_file_path = os.path.join(os.environ.get("SystemRoot", "C:\\Windows"), "inf", "setupapi.dev.log")
    
    # 用于查找指示设备安装的行并捕获供应商ID、产品ID、序列号和时间戳的正则表达式。
    # 示例行：>>>   [Device Install (Hardware IDs): USB\VID_0781&PID_5581\AA010304091047200170 - ROOT\USB\0000]
    device_install_regex = re.compile(r"USB\\VID_([0-9A-F]{4})&PID_([0-9A-F]{4})(?:\\|&)([^\s\]]+)")
    # 用于“Section start”或“Section end”行的时间戳正则表达式
    timestamp_regex = re.compile(r">>>\s*Section (?:start|end)\s*(\d{4}/\d{2}/\d{2}\s*\d{2}:\d{2}:\d{2}(?:\.\d{3})?)")
    
    current_timestamp = None

    try:
        with open(log_file_path, 'r', encoding='utf-8', errors='ignore') as f:
            for line in f:
                ts_match = timestamp_regex.search(line)
                if ts_match:
                    try:
                        # 首先尝试解析带毫秒的时间戳
                        current_timestamp = datetime.datetime.strptime(ts_match.group(1), "%Y/%m/%d %H:%M:%S.%f")
                    except ValueError:
                        # 如果没有毫秒，则尝试不带毫秒的格式
                        try:
                            current_timestamp = datetime.datetime.strptime(ts_match.group(1), "%Y/%m/%d %H:%M:%S")
                        except ValueError:
                            current_timestamp = None # 如果解析出错，则设置为None

                if "Device Install (Hardware IDs)" in line and current_timestamp:
                    match = device_install_regex.search(line)
                    if match:
                        vid = match.group(1).upper()
                        pid = match.group(2).upper()
                        
                        # 序列号可能在'&'后面有额外信息，将其移除
                        serial_full = match.group(3).upper()
                        serial = serial_full.split('&')[0] 

                        device_key = (vid, pid, serial)
                        
                        if device_key in devices_info:
                            # 仅当此时间戳更早（首次安装）时才更新
                            if devices_info[device_key]["first_seen_setupapi"] is None or \
                               current_timestamp < devices_info[device_key]["first_seen_setupapi"]:
                                devices_info[device_key]["first_seen_setupapi"] = current_timestamp
                        else:
                            # 在setupapi日志中找到但（尚未）在注册表扫描中找到的设备。添加它。
                            # 这可能发生在设备已连接但驱动程序未完全枚举，
                            # 或者注册表解析遗漏了它。
                            devices_info[device_key] = {
                                "vid": vid,
                                "pid": pid,
                                "serial_number": serial,
                                "device_type": "从SetupAPI日志", # 标记为主要从日志中找到
                                "description": "N/A (从日志)",
                                "friendly_name": "N/A (从日志)",
                                "manufacturer": "N/A (从日志)",
                                "first_seen_setupapi": current_timestamp,
                                "last_seen_registry": None, # 无法从setupapi日志中获取
                                "registry_path": "N/A (从日志)"
                            }

    except FileNotFoundError:
        print(f"setupapi.dev.log未找到，路径为 {log_file_path}")
    except Exception as e:
        print(f"读取或解析setupapi.dev.log时出错: {e}")


def main():
    print("正在尝试收集USB设备信息。这需要管理员权限才能完全访问。")
    
    # 1. 从注册表获取信息
    usb_devices_data = get_usb_info_from_registry()
    
    # 2. 从setupapi.dev.log获取首次安装时间并合并
    get_first_install_times_from_setupapi(usb_devices_data)

    print(f"\n找到 {len(usb_devices_data)} 个唯一的USB设备条目。\n")

    output_data = []
    # 按照“Last Known (Registry)”时间升序排序。
    # None值将排在最后（通过将None视为True，其他为False来处理）。
    sorted_devices = sorted(
        usb_devices_data.items(),
        key=lambda item: (
            item[1]["last_seen_registry"] is None,  # 将None值排在最后
            item[1]["last_seen_registry"]           # 然后按实际的datetime对象排序
        )
    )

    for key, info in sorted_devices:
        # 确保所有字段都存在，即使为空，以保证CSV的一致性
        row = {
            "VID": info.get("vid", "N/A"),
            "PID": info.get("pid", "N/A"),
            "Serial Number": info.get("serial_number", "N/A"),
            "Device Type": info.get("device_type", "N/A"),
            # 移除可能存在的空字节
            "Description": info.get("description", "").replace("\x00", ""), 
            "Friendly Name": info.get("friendly_name", "").replace("\x00", ""),
            "Manufacturer": info.get("manufacturer", "").replace("\x00", ""),
            "First Seen (SetupAPI)": info["first_seen_setupapi"].strftime('%Y-%m-%d %H:%M:%S') if info["first_seen_setupapi"] else "N/A",
            "Last Known (Registry)": info["last_seen_registry"].strftime('%Y-%m-%d %H:%M:%S') if info["last_seen_registry"] else "N/A",
            "Registry Path": info.get("registry_path", "N/A")
        }
        output_data.append(row)
        
        # 打印到控制台
        print(f"供应商ID (VID): {info['vid']}, 产品ID (PID): {info['pid']}, 序列号: {info['serial_number']}")
        print(f"   类型: {info['device_type']}")
        if info['description']: print(f"   描述: {info['description']}")
        if info['friendly_name']: print(f"   友好名称: {info['friendly_name']}")
        if info.get('manufacturer'): print(f"   制造商: {info.get('manufacturer')}")
        print(f"   首次发现 (SetupAPI): {info['first_seen_setupapi'].strftime('%Y-%m-%d %H:%M:%S') if info['first_seen_setupapi'] else 'N/A'}")
        print(f"   最后已知 (注册表): {info['last_seen_registry'].strftime('%Y-%m-%d %H:%M:%S') if info['last_seen_registry'] else 'N/A'}")
        if info.get('registry_path'): print(f"   注册表路径: {info.get('registry_path')}")
        print("-" * 30)

    # 可选：保存到CSV
    if output_data:
        csv_file = "usb_device_history.csv"
        print(f"\n正在将数据保存到 {csv_file}...")
        try:
            # 获取所有字典中所有可能的字段名，以确保CSV标题完整
            fieldnames = sorted(list(set(k for d in output_data for k in d.keys())))
            with open(csv_file, 'w', newline='', encoding='utf-8') as f:
                writer = csv.DictWriter(f, fieldnames=fieldnames)
                writer.writeheader()
                writer.writerows(output_data)
            print(f"成功保存到 {csv_file}")
        except IOError as e:
            print(f"保存CSV时出错: {e}")

if __name__ == "__main__":
    # 在Windows上可靠地检查管理员权限
    is_admin = False
    if sys.platform == "win32":
        import ctypes
        try:
            is_admin = ctypes.windll.shell32.IsUserAnAdmin() != 0
        except Exception:
            pass # 如果检查失败，则假定不是管理员
    else: # 对于其他操作系统，root通常是0
        try:
            is_admin = (os.geteuid() == 0)
        except AttributeError: # 非Unix-like系统可能没有geteuid
            pass

    if not is_admin:
        print("警告：此脚本可能需要管理员权限才能访问所有注册表项和系统日志。")
        print("如果遇到错误或数据缺失，请尝试以管理员身份运行。")
        # 在Windows上，你可以在这里尝试提升权限
        # 对于一个简单的脚本，仅发出警告通常就足够了。

    main()
